from transformers import pipeline
from transformers import AutoTokenizer,AutoModelForCausalLM

pretrain_model_dir = "/media/dengyunfei/6T/data/models/huggingface/bloom-389m-zh"
# pretrain_model_dir = "/media/dengyunfei/6T/data/logs/run_clm/results/checkpoint-2880"

tokenizer = AutoTokenizer.from_pretrained(pretrain_model_dir)
model = AutoModelForCausalLM.from_pretrained(pretrain_model_dir)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
ipt = "你是谁?"
print(pipe(ipt,
    # max_length=512,
    do_sample=True,
    max_new_tokens=256,
    top_p=0.95,
    temperature=0.7,
    num_beams=2
))